/*
Copyright (c) 2012 Eric Lasota

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

#include <string.h>
#include <stdio.h>
#include <new>

#include "lwmpeg.h"
#include "ringbuffer.h"
#include "../m1vdec/m1vdec.h"
#include "../mp2dec/mp2dec/mp2dec.h"

#define DEMUX_BUFFER_SIZE	10000

static Uint32 ParseUInt32(const void *bytes)
{
	const Uint8 *asBytes = static_cast<const Uint8 *>(bytes);
	return (asBytes[0] << 24) | (asBytes[1] << 16) | (asBytes[2] << 8) | asBytes[3];
}

class LWMPEG_VideoRingBuffer : public LWMPEG_MPEGstream
{
private:
	RingBuffer<M1VDEC_BUFFER_SIZE> _ringBuffer;

public:
	virtual int copy_data(void *destBytes, unsigned int requested)
	{
		return static_cast<int>(_ringBuffer.Read(destBytes, static_cast<size_t>(requested)));
	}

	virtual void peek_bytes(void *destBytes, unsigned int requested)
	{
		_ringBuffer.Peek(destBytes, requested);
	}

	virtual bool eof()
	{
		if(_ringBuffer.BytesStored() == 0)
			return true;
		return false;
	}

	RingBuffer<M1VDEC_BUFFER_SIZE> *RingBuffer()
	{
		return &_ringBuffer;
	}
};

struct LWMPEG_VideoDecodeState
{
private:
	enum
	{
		DemuxTargetAudio,
		DemuxTargetVideo,
		DemuxTargetTrash,
		DemuxTargetParsing,
	};

	enum
	{
		PS_PackHeader,
		PS_SystemHeader,
		PS_SystemHeaderStreams,
		PS_DemuxPacket,
		PS_Exhausted,
	};

	LWMPEG_VideoRingBuffer _videoStream;
	RingBuffer<MP2DEC_MAX_FRAME_SIZE_BYTES + MP2DEC_HEADER_SIZE_BYTES> _audioRingBuffer;
	RingBuffer<DEMUX_BUFFER_SIZE> _demuxBuffer;
	Uint8 _audioData[MP2DEC_MAX_FRAME_SIZE_BYTES];
	Uint8 _demuxMoveBuffer[DEMUX_BUFFER_SIZE];
	mp2dec_sint16 _audioSamples[2][MP2DEC_FRAME_NUM_SAMPLES];
	mp2dec_header _audioHeader;
	mp2dec_decode_properties _audioDec;
	LWMPEG_MPEGvideo *_videoDec;
	LWMPEG_Allocator _alloc;

	bool _parsedSystemHeader;
	bool _flushing;
	bool _audioAvailable;
	bool _videoExhausted;	// If this is true, then no more video will ever be decoded
	bool _mustParsePacketHeaders;

	int _parseState;

	int _demuxTarget;
	bool _noDemuxLimit;
	size_t _demuxRemaining;

public:
	explicit LWMPEG_VideoDecodeState(const LWMPEG_Allocator *alloc)
	{
		mp2dec_init_decode_state(&_audioDec);
		_videoDec = NULL;
		_parsedSystemHeader = false;
		_parseState = PS_PackHeader;
		_flushing = false;
		_demuxTarget = DemuxTargetParsing;
		_demuxRemaining = 0;
		_noDemuxLimit = false;
		_videoExhausted = false;
		_audioAvailable = false;
		_mustParsePacketHeaders = false;

		_alloc = *alloc;

		memset(_audioData, 0, sizeof(_audioData));
	}

	~LWMPEG_VideoDecodeState()
	{
		if(_videoDec)
		{
			_videoDec->~LWMPEG_MPEGvideo();
			_alloc.freefunc(_alloc.opaque, _videoDec);
			_videoDec = NULL;
		}
	}

	void DestroyAndRelease()
	{
		LWMPEG_Allocator alloc = this->_alloc;
		this->~LWMPEG_VideoDecodeState();
		alloc.freefunc(alloc.opaque, this);
	}

	size_t BytesFree() const
	{
		return _demuxBuffer.BytesFree();
	}

	void FeedData(const void *data, size_t size)
	{
		_demuxBuffer.Write(data, size);
	}

	void TryParsePackHeader()
	{
		if(_parseState != PS_PackHeader)
			return;

		if(_demuxBuffer.BytesStored() < 4)
			return;

		Uint8 header[4];
		_demuxBuffer.Peek(header, 4);

		if(ParseUInt32(header) != 0x000001ba)
		{
			_parseState = PS_DemuxPacket;
			_noDemuxLimit = true;

			if(ParseUInt32(header) == 0x000001b3)					// Detect MPEG-1 video start
				_demuxTarget = DemuxTargetVideo;
			else if(header[0] == 0xf && (header[1] & 0xf0) == 0xf0)	// Detect layer 2 sync word
				_demuxTarget = DemuxTargetAudio;
			else
				_demuxTarget = DemuxTargetTrash;

			return;
		}

		// Try to parse the header info
		if(_demuxBuffer.BytesStored() < 24)
			return;

		Uint8 packHeaderBytes[12];
		_demuxBuffer.Read(packHeaderBytes, 12);
		
		// We can just ignore all of this

		// Pack:
		//(b[4] & 0xf0) == 0x20
		//SCR = ((b[4] & 0x7) << 30) | ((b[5]) << 22) | ((b[6] & 0xfe) << 15) | ((b[7]) << 7) | ((b[8] & 0xfe) >> 1);
		//MuxRate = ((b[9] & 0x7f) << 15); | (b[10] << 7) | ((b[11] & 0xfe) >> 1)

		_parseState = PS_DemuxPacket;
	}

	void TryParseSystemHeader()
	{
		if(_parseState != PS_SystemHeader)
			return;

		// This should only ever be reachable by TryDemux setting it
		
		Uint8 systemHeaderBytes[12];
		_demuxBuffer.Read(systemHeaderBytes, 12);
		// System:
		// RateBound = ((b[8] & 0xfe) >> 1) | (b[7] << 7) | ((b[6] & 0x7f) << 15)
		// AudioBound = ((b[9] & 0xfc) >> 2)

		_parseState = PS_SystemHeaderStreams;
	}

	
	void TryParseStreams()
	{
		if(_parseState != PS_SystemHeaderStreams)
			return;

		if(_demuxBuffer.BytesStored() < 3)
			return;
		
		Uint8 marker[3];
		_demuxBuffer.Peek(marker, 3);

		if(marker[0] == 0 && marker[1] == 0 && marker[2] == 1)
		{
			_parseState = PS_DemuxPacket;
			_demuxTarget = DemuxTargetParsing;
		}
		else
			_demuxBuffer.Skip(3);	// Skip past it
	}

	void TryDemux()
	{
		if(_parseState != PS_DemuxPacket)
			return;

		if(_demuxTarget == DemuxTargetParsing)
		{
			if(_demuxBuffer.BytesStored() >= 6)
			{
				Uint8 headerBytes[6];
				_demuxBuffer.Peek(headerBytes, 6);

				if(headerBytes[3] == 0xc0)
					_demuxTarget = DemuxTargetAudio;
				else if(headerBytes[3] == 0xe0)
					_demuxTarget = DemuxTargetVideo;
				else if(ParseUInt32(headerBytes) == 0x1ba)
				{
					// Hit an inline pack header
					_parseState = PS_PackHeader;
					return;
				}
				else if(ParseUInt32(headerBytes) == 0x1bb)
				{
					// Hit an inline system header
					_parseState = PS_SystemHeader;
					return;
				}
				else
					_demuxTarget = DemuxTargetTrash;

				_mustParsePacketHeaders = true;
				_demuxRemaining = (headerBytes[4] << 8) | headerBytes[5];
				_demuxBuffer.Skip(6);
			}
			else if(_flushing)
				_parseState = PS_Exhausted;
		}

		if(_demuxTarget != DemuxTargetParsing)
		{
			if(_mustParsePacketHeaders)
			{
				Uint8 flagBytes[3];
				// Need to skip the packet headers
				if(_demuxBuffer.BytesStored() < 1)
				{
					if(_flushing)
						_parseState = PS_Exhausted;
					return;
				}
				_demuxBuffer.Peek(flagBytes, 1);

				size_t skipSize = 0;
				size_t stdSize = 0;
				size_t ptsdtsSize = 0;
				if(flagBytes[0] != 0x0f)	// Skip STD/PTS/DTS
				{
					if(_demuxBuffer.BytesStored() < 2)
					{
						if(_flushing)
							_parseState = PS_Exhausted;
						return;
					}

					_demuxBuffer.Peek(flagBytes, 2);
					if((flagBytes[0] & 0xc0) == 0x40)
						stdSize = 2;

					if(_demuxBuffer.BytesStored() < 1 + stdSize)
					{
						if(_flushing)
							_parseState = PS_Exhausted;
						return;
					}
					_demuxBuffer.Peek(flagBytes, 1 + stdSize);
					if((flagBytes[stdSize] & 0xf0) == 0x30)
						ptsdtsSize = 10;	// Both PTS and DTS are present
					else if((flagBytes[stdSize] & 0xf0) == 0x20)
						ptsdtsSize = 5;	// Only PTS is present
					
					skipSize = stdSize + ptsdtsSize;
				}
				else
					skipSize = 1;
				if(_demuxBuffer.BytesStored() < skipSize)
				{
					if(_flushing)
						_parseState = PS_Exhausted;
					return;
				}
				_demuxBuffer.Skip(skipSize);
				if(_demuxRemaining < skipSize)
					_demuxRemaining = 0;	// Ugh
				else
					_demuxRemaining -= skipSize;

				_mustParsePacketHeaders = false;
			}

			if(!_mustParsePacketHeaders && _demuxRemaining)
			{
				// See how much data can be moved
				size_t amount = _demuxBuffer.BytesStored();
				if(amount > _demuxRemaining)
					amount = _demuxRemaining;

				if(amount == 0)
				{
					// Nothing left in the demux buffer
					if(_flushing)
						_parseState = PS_Exhausted;
					return;
				}

				if(_demuxTarget == DemuxTargetAudio)
				{
					if(_audioRingBuffer.BytesFree() < amount)
						amount = _audioRingBuffer.BytesFree();

					if(amount)
					{
						_demuxBuffer.Read(_demuxMoveBuffer, amount);
						_audioRingBuffer.Write(_demuxMoveBuffer, amount);
					}
				}
				else if(_demuxTarget == DemuxTargetVideo)
				{
					if(_videoStream.RingBuffer()->BytesFree() < amount)
						amount = _videoStream.RingBuffer()->BytesFree();

					if(amount)
					{
						_demuxBuffer.Read(_demuxMoveBuffer, amount);
						_videoStream.RingBuffer()->Write(_demuxMoveBuffer, amount);
					}
				}
				else if(_demuxTarget == DemuxTargetTrash)
					_demuxBuffer.Skip(amount);

				_demuxRemaining -= amount;
			}

			if(_demuxRemaining == 0)
				_demuxTarget = DemuxTargetParsing;
		}
	}

	void TryDecodeVideo()
	{
		if(_videoExhausted)
		{
			// Throw out any remaining video data
			_videoStream.RingBuffer()->Skip(_videoStream.RingBuffer()->BytesStored());
			return;
		}

		// Require enough preload that this is safe
		if(_videoStream.RingBuffer()->BytesStored() < M1VDEC_BUFFER_SIZE && _parseState != PS_Exhausted)
			return;

		if(!_videoDec)
		{
			_videoDec = static_cast<LWMPEG_MPEGvideo *>(_alloc.allocfunc(_alloc.opaque, sizeof(LWMPEG_MPEGvideo)));
			if(_videoDec == NULL)
				_videoExhausted = true;
			else
				new (_videoDec) LWMPEG_MPEGvideo(&this->_videoStream, 24, &_alloc);
			return;	// Don't do any more parsing this frame
		}
		
		if(!_videoDec->DecodeFrame())
			_videoExhausted = true;
	}

	void TryDecodeAudio()
	{
		// Require enough preload that this is safe
		if(_audioRingBuffer.BytesStored() < MP2DEC_MAX_FRAME_SIZE_BYTES + MP2DEC_HEADER_SIZE_BYTES && _parseState != PS_Exhausted)
			return;

		if(_audioRingBuffer.BytesStored() < MP2DEC_HEADER_SIZE_BYTES)
			return;	// Nothing left

		Uint8 headerBytes[MP2DEC_HEADER_SIZE_BYTES];

		_audioRingBuffer.Read(headerBytes, MP2DEC_HEADER_SIZE_BYTES);
		if(mp2dec_parse_header(headerBytes, &_audioHeader, &_audioDec))
		{
			_audioRingBuffer.Read(_audioData, _audioDec.frameBytes);
			mp2dec_decode_frame(&_audioHeader, &_audioDec, _audioData, _audioSamples);

			_audioAvailable = true;
		}
	}

	bool GetAudio(const short **leftChannel, const short **rightChannel, unsigned int *sampleRate, unsigned int *numChannels, unsigned int *numSamples)
	{
		if(!_audioAvailable)
			return false;

		if(leftChannel) *leftChannel = _audioSamples[0];
		if(rightChannel) *rightChannel = _audioSamples[1];
		if(sampleRate) *sampleRate = _audioDec.sampleRate;
		if(numChannels) *numChannels = (_audioDec.inputStereo ? 2 : 1);
		if(numSamples) *numSamples = MP2DEC_FRAME_NUM_SAMPLES;

		_audioAvailable = false;

		return true;
	}

	bool GetImage(unsigned int *frameRateNum, unsigned int *frameRateDenom, unsigned int *width, unsigned int *height,
		unsigned int *lpitch, unsigned int *cpitch, const unsigned char **yPlane, const unsigned char **uPlane, const unsigned char **vPlane)
	{
		if(!_videoDec)
			return false;
		return _videoDec->GetImage(frameRateNum, frameRateDenom, width, height, lpitch, cpitch, yPlane, uPlane, vPlane);
	}

	void Work()
	{
		TryParsePackHeader();
		TryParseSystemHeader();
		TryParseStreams();
		TryDemux();
		TryDecodeAudio();
		TryDecodeVideo();
	}

	void Flush()
	{
		_flushing = true;
	}

	bool Exhausted()
	{
		return (_parseState == PS_Exhausted) && !_audioAvailable && _videoExhausted;
	}

	void SetFrameSkip(int numFrames)
	{
		if(_videoDec)
			_videoDec->SetFrameSkip(numFrames);
	}
};

LWMPEG_LINKAGE LWMPEG_VideoDecodeState *LWMPEG_CreateDecodeState(const LWMPEG_Allocator *alloc)
{
	LWMPEG_VideoDecodeState *decState = static_cast<LWMPEG_VideoDecodeState *>(alloc->allocfunc(alloc->opaque, sizeof(LWMPEG_VideoDecodeState)));

	if(!decState)
		return NULL;
	new (decState) LWMPEG_VideoDecodeState(alloc);

	return decState;
}

LWMPEG_LINKAGE void LWMPEG_DestroyDecodeState(LWMPEG_VideoDecodeState *state)
{
	state->DestroyAndRelease();
}

LWMPEG_LINKAGE size_t LWMPEG_BytesFree(const LWMPEG_VideoDecodeState *state)
{
	return state->BytesFree();
}

LWMPEG_LINKAGE void LWMPEG_FeedData(LWMPEG_VideoDecodeState *state, const void *data, size_t size)
{
	state->FeedData(data, size);
}

LWMPEG_LINKAGE void LWMPEG_Work(LWMPEG_VideoDecodeState *state)
{
	state->Work();
}

LWMPEG_LINKAGE void LWMPEG_SetFrameSkip(LWMPEG_VideoDecodeState *state, int numFrames)
{
	state->SetFrameSkip(numFrames);
}

LWMPEG_LINKAGE void LWMPEG_Flush(LWMPEG_VideoDecodeState *state)
{
	state->Flush();
}

LWMPEG_LINKAGE int LWMPEG_Exhausted(LWMPEG_VideoDecodeState *state)
{
	return state->Exhausted() ? 1 : 0;
}

LWMPEG_LINKAGE int LWMPEG_GetAudio(LWMPEG_VideoDecodeState *state, const short **leftChannel, const short **rightChannel, unsigned int *sampleRate,
	unsigned int *numChannels, unsigned int *numSamples)
{
	return state->GetAudio(leftChannel, rightChannel, sampleRate, numChannels, numSamples) ? 1 : 0;
}

LWMPEG_LINKAGE int LWMPEG_GetImage(LWMPEG_VideoDecodeState *state, unsigned int *frameRateNum, unsigned int *frameRateDenom,
	unsigned int *width, unsigned int *height, unsigned int *lpitch, unsigned int *cpitch,
	const unsigned char **yPlane, const unsigned char **uPlane, const unsigned char **vPlane)
{
	return state->GetImage(frameRateNum, frameRateDenom, width, height, lpitch, cpitch, yPlane, uPlane, vPlane) ? 1 : 0;
}
